package com.starinside.wp.listener;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.unfbx.chatgpt.entity.chat.ChatChoice;
import com.unfbx.chatgpt.entity.chat.ChatCompletionResponse;
import com.unfbx.chatgpt.entity.chat.Message;
import com.zhipu.oapi.service.v3.StandardEventSourceListener;
import lombok.Data;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Response;
import okhttp3.sse.EventSource;
import org.apache.commons.lang3.StringUtils;

import java.util.Arrays;
import java.util.List;

/**
 * 描述：OpenAIEventSourceListener
 *
 * @author https:www.unfbx.com
 * @date 2023-02-22
 */
@Slf4j
@Data
public class MockEventSourceListener extends StandardEventSourceListener  {

    private StringBuffer streamCache = new StringBuffer();

    private OpenAIEventSourceListener mockListener;

    public MockEventSourceListener(OpenAIEventSourceListener mockListener){
        this.mockListener = mockListener;
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void onOpen(EventSource eventSource, Response response) {

        super.onOpen(eventSource, response);

        mockListener.onOpen(eventSource, response);

    }

    /**
     * {@inheritDoc}
     */
    @SneakyThrows
    @Override
    public void onEvent(EventSource eventSource, String id, String type, String data) {
//        log.info("MOCK返回数据：{}", data);

        super.onEvent(eventSource, id, type, data);

        if ("finish".equals(type)) {

            if (streamCache.length() > 0) {//clear buffer
                mockOnEvent(eventSource, id, type, streamCache.toString());
                streamCache = new StringBuffer();
            }

            mockListener.onEvent(eventSource, id, type, "[DONE]");
            return;
        }

        if (StringUtils.isBlank(data)){
            return;
        }

        JSONObject jsonObject = JSON.parseObject(data);
        String curText = jsonObject.getString("data");
        streamCache.append(curText);  //append to buffer

        if (streamCache.length() > 5) { // use and clear buffer
            mockOnEvent(eventSource, id, type, streamCache.toString());
            streamCache = new StringBuffer();
        }
    }

    private void mockOnEvent(EventSource eventSource, String id, String type, String curText){
        ChatChoice choice = new ChatChoice();
        choice.setDelta(Message.builder().role(Message.Role.ASSISTANT).content(curText).build());

        List<ChatChoice> choices = Arrays.asList(choice);

        ChatCompletionResponse completionResponse = new ChatCompletionResponse();
        completionResponse.setChoices(choices);
        completionResponse.setId("" + System.currentTimeMillis());

        mockListener.onEvent(eventSource, id, type, JSON.toJSONString(completionResponse));
    }


    @Override
    public void onClosed(EventSource eventSource) {
        log.info("MOCK关闭链接");

        super.onClosed(eventSource);

        mockListener.onClosed(eventSource);
    }


    @SneakyThrows
    @Override
    public void onFailure(EventSource eventSource, Throwable t, Response response) {
        log.info("MOCK 失败");
        t.printStackTrace();

        super.onFailure(eventSource, t, response);

        mockListener.onFailure(eventSource, t, response);
    }
}
